use std::sync::Arc;

use async_trait::async_trait;
use rig::{agent::{Agent, AgentBuilder}, completion::Prompt, providers::deepseek::{self, DeepSeekCompletionModel}};

use common::compose::{CallBox, CallCore, ContextBox, FError};

pub struct AgentConfig {
    pub temperature: f32,
    pub max_tokens: usize,
    pub system: String,
}

pub fn init_agent<'a>(system: &'a str, api_key: &'a str) -> CallBox<'a, String, String> {
    let client = deepseek::Client::new(api_key);
    let model = client.completion_model(deepseek::DEEPSEEK_CHAT);
    let agent = AgentBuilder::new(model)
        .temperature(1.3)
        .preamble(system)
        .max_tokens(8192)
        .build();

    struct AgentWrapper {
        agent: Agent<DeepSeekCompletionModel>,
    }

    #[async_trait]
    impl CallCore<String, String> for AgentWrapper {
        async fn run(&self, prompt: String, _: ContextBox) -> Result<String, FError> {
            let chat_response = self.agent.prompt(prompt)
                .await
                .expect("Failed to prompt the Agent");
            Ok(chat_response)
        }  
    }

    Arc::new(Box::new(AgentWrapper { agent }))
}


#[cfg(test)]
mod tests {
    use common::compose::make_empty_ctx;

    use super::*;
    use std::env;
    use tokio::runtime::Runtime;

    #[test]
    fn test_init_agent() {
        dotenvy::from_path("../../.env").ok();
        // 从环境变量获取API密钥，如果没有设置则跳过测试
        let api_key = match env::var("DEEPSEEK_API_KEY") {
            Ok(key) => key,
            Err(_) => {
                println!("跳过测试：未设置DEEPSEEK_API_KEY环境变量");
                return;
            }
        };

        // 创建系统提示
        let system = r#"
你是一个HTML结构化提取专家，请将输入的HTML内容转换为符合规范的JSON目录树。规则如下：
# 处理规则
1. 每个节点必须包含title字段，取自最近a标签的文本内容（若无a标签则取父级文本）
2. 仅当存在有效链接时添加url字段，取自对应a标签的href属性
3. 嵌套结构处理：
   - 当遇到<ul>包裹的<li>列表时，生成children数组
   - 子节点必须比父节点多一级缩进（通过DOM层级判断）
   - 递归处理嵌套结构，最多支持5级嵌套
4. 请直接返回json字符串，不要markdown格式

# 输出格式
[
  {
    "title": "章节标题",
    "url?": "可选链接",
    "children?": [子节点]
  },
  ...
]

# 示例演示
输入HTML：
<div class="toc">
  <ul>
    <li><a href="/intro">概述</a></li>
    <li>产品特性
      <ul>
        <li><a href="/core">核心功能</a></li>
        <li><a href="/ui">交互设计</a></li>
      </ul>
    </li>
  </ul>
</div>

期望输出：
[
  {
    "title": "概述",
    "url": "/intro"
  },
  {
    "title": "产品特性",
    "children": [
      {
        "title": "核心功能",
        "url": "/core"
      },
      {
        "title": "交互设计", 
        "url": "/ui"
      }
    ]
  }
]

# 特殊场景处理
1. 混合层级：
输入：<li>第一章<ul>...</ul></li>
输出：父级title取自"第一章"，子级从<ul>提取

2. 无链接节点：
输入：<li><span>附录</span></li>
输出：{"title": "附录"}

3. 深层嵌套：
输入：li>ul>li>ul>li结构对应生成3级children数组
"#;
        
        // 初始化Agent
        let agent = init_agent(system, &api_key);
        
        // 创建Tokio运行时来执行异步代码
        let rt = Runtime::new().unwrap();
        
        // 执行异步测试
        let result = rt.block_on(async {
            // 发送测试提示
            let prompt = r#"test"#;
            let response = agent.run(prompt.to_string(), make_empty_ctx()).await;
            
            // 验证响应
            assert!(response.is_ok(), "Agent应该返回一个有效的响应");
            let response_text = response.unwrap();
            
            // 验证响应不为空
            assert!(!response_text.is_empty(), "Agent的响应不应该为空");
            
            // 打印响应以便手动检查
            println!("Agent响应: {}", response_text);
            
            Ok::<(), String>(())
        });
        
        // 确保测试成功完成
        assert!(result.is_ok());
    }

  }